import torch
from torch import nn


class FishNet(nn.Sequential):
    def __init__(self, in_ch, out_ch):
        layers = [
            nn.Linear(in_ch, 16),
            nn.LeakyReLU(),
            nn.Linear(16, out_ch)
        ]
        # 定义好层级后才调用构造函数，否则不能通过父类的方法正确载入训练结果
        super(FishNet, self).__init__(*layers)
        # 将weight_init初始化方法应用到模块上，会自动递归调用
        self.apply(weight_init)


def weight_init(m):
    # 递归调用初始化所有的线性连接层符合正态分布
    if isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, 0, 0.1)
        if m.bias is not None:
            nn.init.constant_(m.bias, 0)
